#用于把labelimg导出的默认xml格式转成yolo可用的txt格式
import os
import pickle
import xml.etree.ElementTree as ET
from os import listdir, getcwd
from os.path import join


def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]

    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]

    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            # b=(xmin, xmax, ymin, ymax)
            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    # 把voc的xml标签文件转化为yolo的txt标签文件
    # 1、类别
    classes1 = ['钩翅眼蛱蝶', '老豹蛱蝶', '红基美凤蝶', '锦瑟蛱蝶', '伊诺小豹蛱蝶', '隐纹谷弄蝶', '红襟粉蝶', '柳紫闪蛱蝶', '大紫琉璃灰蝶', '链环蛱蝶', '金裳凤蝶', '美眼蛱蝶',
                '青凤蝶', '突角小粉蝶', '波太玄灰蝶', '亮灰蝶', '灿福蛱蝶', '翠蓝眼蛱蝶', '荨麻蛱蝶', '尖翅翠蛱蝶', '雅弄蝶', '柱菲蛱蝶', '中环蛱蝶', '阿芬眼蝶', '虎斑蝶',
                '黎明豆粉蝶', '曲纹紫灰蝶', '云粉蝶', '银斑豹蛱蝶', '维纳斯眼灰蝶', '艳灰蝶', '咖灰蝶', '菜粉蝶', '琉璃蛱蝶', '无斑珂弄蝶', '密纹飒弄蝶', '青海红珠灰蝶',
                '蟾福蛱蝶', '线灰蝶', '大卫粉蝶', '曲斑珠蛱蝶', '绢斑蝶', '依帕绢蝶', '巴黎翠凤蝶', '白眼蝶', '白钩蛱蝶', '大翅绢粉蝶', '柑橘凤蝶', '绿豹蛱蝶', '边纹黛眼蝶',
                '碧凤蝶', '翠袖锯眼蝶', '黑网蛱蝶', '扬眉线蛱蝶', '绢蛱蝶', '蓝凤蝶', '银豹蛱蝶', '蓝灰蝶', '镉黄迁粉蝶', '宽边黄粉蝶', '秀蛱蝶', '珍蛱蝶', '侏粉蝶',
                '箭纹云粉蝶', '婀灰蝶', '朴喙蝶', '素弄蝶', '玄珠带蛱蝶', '牧女珍眼蝶', '斐豹蛱蝶', '蛇目褐蚬蝶', '云豹蛱蝶', '直纹稻弄蝶', '珍珠绢蝶', '玄灰蝶',
                '蓝点紫斑蝶', '红灰蝶', '古北拟酒眼蝶', '西门珍眼蝶', '四川绢蝶', '虬眉带蛱蝶', '拟稻眉眼蝶', '橙黄豆粉蝶', '山豆粉蝶', '网蛱蝶', '菩萨酒眼蝶', '黄钩蛱蝶',
                '玉带凤蝶', '花弄蝶', '小红蛱蝶', '箭纹绢粉蝶', '绢粉蝶', '黄环蛱蝶', '小黄斑弄蝶']
    # 2、voc格式的xml标签文件路径
    xml_files1 = r'D:\python\PatternRecognition\TrainData1\labels'
    # xml_files1 = r'Annotations'
    # 3、转化为yolo格式的txt标签文件存储路径
    save_txt_files1 = r'2'
    convert_annotation(xml_files1, save_txt_files1, classes1)
